from typing import Union, Dict

import neurokit2 as nk
import numpy as np


class DenoiseECG(object):
    def __init__(self, sampling_rate: int = 500, method: str = 'neurokit', keys = None) -> None:
        self.sampling_rate = sampling_rate
        self.method = method
        self.keys = keys

    def __call__(self, data: Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]) -> Union[np.ndarray, Dict[str, Union[np.ndarray, float]]]:
        if isinstance(data, dict):
            for key in self.keys:
                sample = data[key]
                
                for lead in range(len(sample)):
                    sample[lead] = nk.ecg_clean(ecg_signal=sample[lead], sampling_rate=self.sampling_rate, method=self.method)

                data[key] = sample

        elif isinstance(data, np.ndarray):
            for lead in range(len(data)):
                data[lead] = nk.ecg_clean(ecg_signal=data[lead], sampling_rate=self.sampling_rate, method=self.method)

        return data